# 部署使用，接口代码

import torch
import flask
from builder import MoCo
from PIL import Image
from torchvision import transforms,models
import torch.nn as nn
import yaml
import cv2
import os
import torch.nn.functional as F 
import requests
from io import BytesIO
# 两张图片进行预测

def load_image(url):
    try:
        headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3',
            'Referer': 'https://pjwv.bmwwx.cn',
            'Accept': 'image/jpeg,image/png,image/svg+xml,image/*;q=0.8',
            'Accept-Language': 'en-US,en;q=0.9',
            'Accept-Encoding': 'gzip, deflate, br'
            }
           
        response = requests.get(url, headers=headers)
        response.raise_for_status()
        content_type = response.headers.get('Content-Type')
        if 'image' not in content_type:
            print(f"Error: {url} 返回的内容类型不是图像，实际类型为 {content_type}")
            return None
        img = Image.open(BytesIO(response.content))
        return img
    except requests.RequestException as e:
        print(f"Error fetching image from {url}: {e}")
        return None
    except Exception as e:
        print(f"Unexpected error processing {url}: {e}")
        return None
        
def get_similarity(model, img_path1, img_path2, device='cpu'):
    
    """
    获取两图相似度
    参数：
        model : 加载好的模型
        img_path1 : 第一张图片路径
        img_path2 : 第二张图片路径
        device : 计算设备
    返回：
        similarity_score : 相似度分数 (0-1)
    """
    # 确定性图像预处理
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        normalize,
    ])
    cos_sim = nn.CosineSimilarity(dim=1)
    # torch.save(model.state_dict(),'model_linear.pth')
    # 保证文件读取的确定性
    img1 = load_image(img_path1)
    img2 = load_image(img_path2)
    if img1 is None or img2 is None:
        return 0
    # 转换张量
    tensor1 = transform(img1.convert('RGB')).unsqueeze(0).to(device)  # [1,3,224,224]
    tensor2 = transform(img2.convert('RGB')).unsqueeze(0).to(device)

    feat1 = model(tensor1)
    feat2 = model(tensor2)
    feat1 = F.normalize(feat1,dim=1)
    feat2 = F.normalize(feat2,dim=1)
    similarity = cos_sim(feat1,feat2)
    # 模型推理
    print(similarity)
    return similarity.item()
    # with torch.no_grad(), Image.open(img_path1) as img1, Image.open(img_path2) as img2:
    #     if img1 is None or img2 is None:
    #         return 0
    #     # 转换张量
    #     tensor1 = transform(img1.convert('RGB')).unsqueeze(0).to(device)  # [1,3,224,224]
    #     tensor2 = transform(img2.convert('RGB')).unsqueeze(0).to(device)

    #     feat1 = model(tensor1)
    #     feat2 = model(tensor2)
    #     feat1 = F.normalize(feat1,dim=1)
    #     feat2 = F.normalize(feat2,dim=1)
    #     similarity = cos_sim(feat1,feat2)
    #     # 模型推理
    #     print(similarity)
    #     return similarity.item()
        
app = flask.Flask(__name__)
device = None
model = None
config = None
# 加载模型
def load_model():
    global model
    # model = MoCoWithSimilarity()
    # model.load_state_dict(torch.load('model_linear.pth'))
       # 加载 MoCo encoder_q
    arch = config['arch']
    base_enc = models.__dict__[arch]
    model_s = MoCo(base_enc, dim=config['moco-dim'], K=0, m=0, T=0)
    # 不需要加载 MoCo checkpoint, encoder_q 结构已足够
    model_s.load_state_dict(torch.load(config['weight_path'],map_location = device))
    model = model_s.encoder_q.to(device)
    model.eval()
    return model
#  读取yaml文件
def load_config(file_path):
    global  config
    global device 
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    try:
        with open(file_path, 'r',encoding='utf-8') as file:
            config = yaml.safe_load(file)  # 安全解析YAML
            return config
    except FileNotFoundError:
        print(f"错误：配置文件 {file_path} 不存在")
        return None
    except yaml.YAMLError as e:
        print(f"YAML解析错误：{e}")
        return None
# 对图片中的名字进行划分
def split_pic_names(image_path,output_dir):
    os.makedirs(output_dir, exist_ok=True)
    # 读取图像
    img = cv2.imread(image_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # **1. 预处理：二值化**
    thresh = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                cv2.THRESH_BINARY_INV, 15, 5)
    # **3. 水平和竖直裁剪：**
    height, width = gray.shape
    # 竖直裁剪范围：0.23 到 0.86
    y_start = int(height * config['height_start'])  # 从竖直方向23%开始
    y_end = int(height *config['height_end'])    # 竖直方向86%结束
    # 水平裁剪范围：0.83 到 0.91
    x_start = int(width * config['width_start'])   # 从水平83%开始
    x_end = int(width * config['width_end'])     # 水平91%结束
    # 先对图像进行水平和竖直裁剪
    cropped_img = thresh[y_start:y_end, x_start:x_end]  # 裁剪二值化图像
    cropped_color_img = img[y_start:y_end, x_start:x_end]  # 裁剪彩色图像
    # 计算分割成20份的每一份大小
    
    slice_height = cropped_img.shape[0] // config['num_slices']  # 每一份的高度
    # **4. 分割裁剪区域：**
    for i in range(config['num_slices']):
        # 竖直方向分割
        current_y_start = i * slice_height
        current_y_end = (i + 1) * slice_height if i < config['num_slices'] - 1 else cropped_img.shape[0]  # 最后一块直到图片底部
        # 提取当前的裁剪区域
        extracted_region = cropped_img[current_y_start:current_y_end, :]
        color_extracted_region = cropped_color_img[current_y_start:current_y_end, :]
        color_output_path = os.path.join(output_dir, f"region_{i+1}.jpg")
        cv2.imwrite(color_output_path, color_extracted_region)  # 保存提取区域的彩色图像
    print("签名区域已保存到：{}".format(config['output_dir']))




# 从后端传入对应的一张图片和生成的图片保存的路径，对图片进行分割，分割出存在对应姓名，并保存到相应的路径中
@app.route('/split_pics',methods=[ 'POST'])
def split_pics():
    data= {
        'success':False,
        'message':''
    }
    if flask.request.method == "POST":
        datas = flask.request.get_json(force=True,silent=True)
        img = datas['image']
        output_dir = datas['output_dir']
        split_pic_names(img,config['output_dir'])
    data['success'] = True
    data['message'] = "图片已进行分割,图片保存到{}".format(config['output_dir'])
    return data

#两张图片进行预测
@app.route('/predict_2',methods=["POST"])
def predict_2():
    data ={
        "success":False,
        'similarity':0,
        'judge':'阈值为{},预测不是同一个人'.format(config['threshold'])
    }
    if flask.request.method == "POST":
        imgs = flask.request.get_json(force=True,silent=True)
        print(imgs['images'])      
        similarity = get_similarity(model, imgs['images'][0], imgs['images'][1],device)
        # print(sim)
        print(similarity) 
        data['similarity'] = similarity
        if similarity > config['threshold']: # 是否为同一个人写的，阈值默认为0.65
            data['judge'] = '阈值为{},预测为预测为同一个人'.format(config['threshold'])
        data['success'] = True
        print(data)
    return data

if __name__ == '__main__':
    load_config('config.yaml')
    load_model()
    app.run(host=config['host'],port=config['port'])